package heap

import (
	"compare"
	"util"
)

const InitCap = 16

type Element interface{}

/*
最小堆
*/
type Heap struct {
	vals []Element
	cmp  compare.Comparator
}

//---公有方法---//
func New() *Heap {
	h := new(Heap)
	h.vals = make([]Element, 0, InitCap)
	return h
}

func Heapify(arr []interface{}) *Heap {
	//利用已有的数组构建最小堆，时间复杂度O(n)
	h := new(Heap)
	h.vals = make([]Element, len(arr), cap(arr))
	for i, e := range arr {
		//初始化比较器
		if h.cmp == nil {
			h.cmp = compare.GetCmp(e)
		}
		//初始化元素
		h.vals[i] = e
	}
	if h.Size() <= 1 {
		return h
	}
	//从最后一个非叶子节点开始逐个下沉元素
	n := len(h.vals)
	for i := parent(n - 1); i >= 0; i-- {
		h.shiftDown(i)
	}
	return h
}

func (this *Heap) SetComparator(cmp compare.Comparator) {
	this.cmp = cmp
}

func (this *Heap) Size() int {
	return len(this.vals)
}

func (this *Heap) Empty() bool {
	return len(this.vals) == 0
}

func (this *Heap) Clear() {
	this.vals = make([]Element, 0, InitCap)
}

func (this *Heap) Push(e Element) {
	//如果比较器未初始化，则初始化
	if this.cmp == nil {
		this.cmp = compare.GetCmp(e)
	}
	//新元素加到末尾
	this.vals = append(this.vals, e)
	//末尾的新元素上浮
	lastIndex := len(this.vals) - 1
	this.shiftUp(lastIndex)
}

func (this *Heap) Pop() Element {
	if this.Empty() {
		panic("heap is empty")
	}
	//交换堆顶与末尾元素
	top := this.vals[0]
	lastIndex := len(this.vals) - 1
	this.swap(0, lastIndex)
	//删除末尾元素
	this.vals = this.vals[:lastIndex]
	//堆顶元素下沉
	this.shiftDown(0)
	return top
}

func (this *Heap) Top() Element {
	if this.Empty() {
		return nil
	}
	return this.vals[0]
}

func (this *Heap) Replace(e Element) Element {
	if this.Empty() {
		panic("heap is empty")
	}
	//一次O(logn)操作弹出堆顶并加入一个新元素
	top := this.vals[0]
	this.vals[0] = e
	this.shiftDown(0)
	return top
}

func (this *Heap) String() string {
	vals2 := make([]interface{}, len(this.vals))
	for i, val := range this.vals {
		vals2[i] = val
	}
	return "Top" + util.ToString(vals2)
}

/*
       1(0)
   /          \
 2(1)        4(2)
/     \      /
3(3)  5(4)  6(5)
*/
//---私有方法---//
func parent(i int) int {
	if i == 0 {
		panic("root has no parent")
	}
	return (i - 1) / 2
}

func leftChild(i int) int {
	return 2*i + 1
}

func rightChild(i int) int {
	return 2*i + 2
}

func (this *Heap) swap(i, j int) {
	vals := this.vals
	vals[i], vals[j] = vals[j], vals[i]
}

func (this *Heap) shiftUp(i int) {
	vals := this.vals
	for i > 0 && this.cmp(vals[parent(i)], vals[i]) > 0 {
		this.swap(parent(i), i)
		i = parent(i)
	}
}

func (this *Heap) shiftDown(i int) {
	vals := this.vals
	n := len(vals)
	for leftChild(i) < n {
		//先取左子节点
		childIndex := leftChild(i)
		child := vals[childIndex]
		if rightChild(i) < n {
			//如果有右子节点，取两个子节点中最小的
			rightChildIndex := rightChild(i)
			rightChild := vals[rightChildIndex]
			if this.cmp(rightChild, child) < 0 {
				child = rightChild
				childIndex = rightChildIndex
			}
		}
		//如果子节点比父节点小，就交换
		if this.cmp(child, vals[i]) < 0 {
			this.swap(childIndex, i)
			i = childIndex
		} else {
			break
		}
	}
}
